{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "教程文档：https://tvm.apache.ac.cn/docs/get_started/tutorials/quick_start.html\n",
        "\n",
        "conda_env:tvm-build"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n",
        "\n",
        "# Quick Start\n",
        "\n",
        "This tutorial is for people who are new to Apache TVM. Taking an simple example\n",
        "to show how to use Apache TVM to compile a simple neural network.\n",
        "    :depth: 2\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Overview\n",
        "Apache TVM is a machine learning compilation framework, following the principle of\n",
        "**Python-first development** and **universal deployment**. It takes in pre-trained\n",
        "machine learning models, compiles and generates deployable modules that can be embedded\n",
        "and run everywhere.\n",
        "Apache TVM also enables customizing optimization processes to introduce new optimizations,\n",
        "libraries, codegen and more.\n",
        "\n",
        "Apache TVM can help to:\n",
        "\n",
        "- **Optimize** performance of ML workloads, composing libraries and codegen.\n",
        "- **Deploy** ML workloads to a diverse set of new environments, including new runtime and new\n",
        "  hardware.\n",
        "- **Continuously improve and customize** ML deployment pipeline in Python by quickly customizing\n",
        "  library dispatching, bringing in customized operators and code generation.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Overall Flow\n",
        "Then we will show the overall flow of using Apache TVM to compile a neural network model,\n",
        "showing how to optimize, deploy and run the model.\n",
        "The overall flow is illustrated as the figure:\n",
        "\n",
        "The overall flow consists of the following steps:\n",
        "\n",
        "- **Construct or Import a Model**: Construct a neural network model or import a pre-trained\n",
        "  model from other frameworks (e.g. PyTorch, ONNX), and create the TVM IRModule, which contains\n",
        "  all the information needed for compilation, including high-level Relax functions for\n",
        "  computational graph, and low-level TensorIR functions for tensor program.\n",
        "- **Perform Composable Optimizations**: Perform a series of optimization transformations,\n",
        "  such as graph optimizations, tensor program optimizations, and library dispatching.\n",
        "- **Build and Universal Deployment**: Build the optimized model to a deployable module to the\n",
        "  universal runtime, and execute it on different devices, such as CPU, GPU, or other accelerators.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Construct or Import a Model\n",
        "Before we get started, let's construct a neural network model first.\n",
        "In this tutorial, to make things simple, we will defined a two-layer MLP networks\n",
        "directly in this script with TVM Relax frontend, which is a similar API to PyTorch.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import tvm\n",
        "from tvm import relax\n",
        "from tvm.relax.frontend import nn\n",
        "\n",
        "\n",
        "class MLPModel(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(MLPModel, self).__init__()\n",
        "        self.fc1 = nn.Linear(784, 256)\n",
        "        self.relu1 = nn.ReLU()\n",
        "        self.fc2 = nn.Linear(256, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.fc1(x)\n",
        "        x = self.relu1(x)\n",
        "        x = self.fc2(x)\n",
        "        return x"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Then we can export the model to TVM IRModule, which is the central intermediate representation\n",
        "in TVM.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
              "\n",
              "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">forward</span>(x: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_weight: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_bias: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">256</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_weight: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_bias: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
              "        R<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;num_input&quot;</span>: <span style=\"color: #008000\">1</span>})\n",
              "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>dataflow():\n",
              "            permute_dims: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">784</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>permute_dims(fc1_weight, axes<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
              "            matmul: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>matmul(x, permute_dims, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;void&quot;</span>)\n",
              "            add: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>add(matmul, fc1_bias)\n",
              "            relu: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>nn<span style=\"color: #A2F; font-weight: bold\">.</span>relu(add)\n",
              "            permute_dims1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>permute_dims(fc2_weight, axes<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
              "            matmul1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>matmul(relu, permute_dims1, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;void&quot;</span>)\n",
              "            add1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>add(matmul1, fc2_bias)\n",
              "            gv: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> add1\n",
              "            R<span style=\"color: #A2F; font-weight: bold\">.</span>output(gv)\n",
              "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "mod, param_spec = MLPModel().export_tvm(\n",
        "    spec={\"forward\": {\"x\": nn.spec.Tensor((1, 784), \"float32\")}}\n",
        ")\n",
        "mod.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Perform Optimization Transformations\n",
        "Apache TVM leverage ``pipeline`` to transform and optimize program.\n",
        "The pipeline encapsulates a collection of transformation that gets two goals (at the same level):\n",
        "\n",
        "- **Model optimizations**: such as operator fusion, layout rewrites.\n",
        "- **Tensor program optimization**: Map the operators to low-level implementations\n",
        "  (both library or codegen)\n",
        "\n",
        "<div class=\"alert alert-info\"><h4>Note</h4><p>The twos are goals but not the stages of the pipeline. The two optimizations are performed\n",
        "  **at the same level**, or separately in two stages.</p></div>\n",
        "\n",
        "<div class=\"alert alert-info\"><h4>Note</h4><p>In this tutorial we only demonstrate the overall flow, by leverage ``zero`` optimization\n",
        "  pipeline, instead of optimizing for any specific target.</p></div>\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "mod = relax.get_pipeline(\"zero\")(mod)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Build and Universal Deployment\n",
        "After the optimization, we can build the model to a deployable module and run it on\n",
        "different devices.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[[25123.062 26992.596 25713.104 24646.633 25388.85  24765.168 25504.607\n",
            "  25699.818 25522.967 24999.217]]\n"
          ]
        }
      ],
      "source": [
        "import numpy as np\n",
        "\n",
        "target = tvm.target.Target(\"llvm\")\n",
        "# ex = tvm.compile(mod, target)\n",
        "ex = relax.build(mod, target)\n",
        "device = tvm.cpu()\n",
        "vm = relax.VirtualMachine(ex, device)\n",
        "data = np.random.rand(1, 784).astype(\"float32\")\n",
        "tvm_data = tvm.nd.array(data, device=device)\n",
        "params = [np.random.rand(*param.shape).astype(\"float32\") for _, param in param_spec]\n",
        "params = [tvm.nd.array(param, device=device) for param in params]\n",
        "print(vm[\"forward\"](tvm_data, *params).numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Our goal is to bring machine learning to the application with any language of interest,\n",
        "with the minimum runtime support.\n",
        "\n",
        "- Each function in IRModule becomes a runnable function in the runtime. For example in LLM\n",
        "  cases, we can call ``prefill`` and ``decode`` functions directly.\n",
        "\n",
        "```Python\n",
        "prefill_logits = vm[\"prefill\"](inputs, weight, kv_cache)\n",
        "decoded_logits = vm[\"decode\"](inputs, weight, kv_cache)\n",
        "```\n",
        "- TVM runtime comes with native data structures, such as NDArray, can also have zero\n",
        "  copy exchange with existing ecosystem (DLPack exchange with PyTorch)\n",
        "\n",
        "```Python\n",
        "# Convert PyTorch tensor to TVM NDArray\n",
        "x_tvm = tvm.nd.from_dlpack(x_torch.to_dlpack())\n",
        "# Convert TVM NDArray to PyTorch tensor\n",
        "x_torch = torch.from_dlpack(x_tvm.to_dlpack())\n",
        "```\n",
        "- TVM runtime works in non-python environments, so it works on settings such as mobile\n",
        "\n",
        "```C++\n",
        "// C++ snippet\n",
        "runtime::Module vm = ex.GetFunction(\"load_executable\")();\n",
        "vm.GetFunction(\"init\")(...);\n",
        "NDArray out = vm.GetFunction(\"prefill\")(data, weight, kv_cache);\n",
        "```\n",
        "```Java\n",
        "// Java snippet\n",
        "Module vm = ex.getFunction(\"load_executable\").invoke();\n",
        "vm.getFunction(\"init\").pushArg(...).invoke;\n",
        "NDArray out = vm.getFunction(\"prefill\").pushArg(data).pushArg(weight).pushArg(kv_cache).invoke();\n",
        "```\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Read next\n",
        "This tutorial demonstrates the overall flow of using Apache TVM to compile a neural network model.\n",
        "For more advanced or specific topics, please refer to the following tutorials\n",
        "\n",
        "\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "tvm-build",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
